package com.wutoon.etl.handler;

import com.wutoon.etl.aop.PowerJobLogReport;
import com.wutoon.etl.constant.CommonConstant;
import com.wutoon.etl.pojo.DataSourceDO;
import com.wutoon.etl.pojo.DataTransferTask;
import com.wutoon.etl.util.DateUtil;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.context.request.RequestContextHolder;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicMarkableReference;

/**
 * @author 武兴云/72176468
 * @version 1.0
 * @date 2024/1/22 17:38
 * @description 数据ETL
 */
@Setter
@Getter
@Slf4j
public class DataTransferHandler {
    private DataTransferTask task;
    /**
     * 阻塞队列
     */
    private BlockingQueue<Object[]> queue = new LinkedBlockingQueue<>(CommonConstant.MAX_FETCH_SIZE);

    private final ConcurrentHashMap<Integer, String> columnMap = new ConcurrentHashMap<>();

    public DataTransferHandler(DataTransferTask task) {
        this.task = task;
    }

    private List<Connection> targetConnections = new ArrayList<>();

    /**
     * 速度测试5W/S
     * 支持全量，增量，清洗
     * 参数示例：
     * {
     *     "shardNum": 1,
     *     "customParam": {
     *         "sourceCode": "lwms",
     *         "querySql": "select BIGCUSTOMERCODE, WHID, WHNAME, LOGITICSCODE, LOGITICSNAME, LOGITICSWHCODE, cast(CREATEDTIME as date), cast(UPDATETIME as date) FROM VIEW_WMS_ERP_WH ",
     *         "targetCode": "task_query_center",
     *         "increase": "false",
     *         "targetTable": "ods_VIEW_WMS_ERP_WH",
     *         "targetSql": "insert into ods_VIEW_WMS_ERP_WH(bigcustomercode, whid, whname, logiticscode, logiticsname, logiticswhcode, createdtime, updatetime) values(?,?,?,?,?,?,?,?)",
     *         "fetchSize":5000,
     *         "batchInsertSize":500,
     *         "concurrencySize":1
     *     }
     * }
     *
     * @author 武兴云/72176468
     * @date 2024/2/29 15:43
     * @update_by 武兴云/72176468
     * @update_at 2024/2/29 15:43
     * @creed if you have doubt , please contact me !!!
     */
    public void start() throws Exception {
        PowerJobLogReport.report("DataTransferHandler start");
        try (Connection sourceConn = connect(task.getDataSource());) {
            // 部分参数做limit调整，例如：fetchSize不能超过2W
            limitParam(task);

            // 1、检查目标sql，白名单控制：只允许执行insert语句
            if (!task.getTargetSql().toLowerCase(Locale.ROOT).trim().startsWith(CommonConstant.START_WITH_INSERT) &&
                    !task.getTargetSql().toLowerCase(Locale.ROOT).trim().startsWith(CommonConstant.START_WITH_UPSERT) &&
                    !task.getTargetSql().toLowerCase(Locale.ROOT).trim().startsWith(CommonConstant.START_WITH_REPLACE_INTO)) {
                throw new RuntimeException("invalid danger target sql:" + task.getTargetSql());
            }

            // 2、初始化连接池，离线任务用得少，连接池实时建立
            for (int i = 0; i < task.getConcurrencySize(); i++) {
                targetConnections.add(connect(task.getTargetSource()));
            }
            // 需要账号具有drop权限，否则会失败
            executePreSQl(targetConnections, task);

            // 3、构建生产者
            ExecutorService executor = Executors.newFixedThreadPool(task.getConcurrencySize() + 1);
            CountDownLatch produceEnd = new CountDownLatch(CommonConstant.NUMBER_ONE);
            CountDownLatch consumerEnd = new CountDownLatch(task.getConcurrencySize());
            // 主线程和子线程之间消息传递
            AtomicMarkableReference<String> msg = new AtomicMarkableReference<>(null, false);
            JdbcReader reader = new JdbcReader(queue, columnMap, task, sourceConn, produceEnd, msg, RequestContextHolder.getRequestAttributes());
            executor.execute(reader);

            // 4、构建消费者
            for(int i = 0; i < task.getConcurrencySize(); i ++) {
                JdbcWriter writer = new JdbcWriter(queue, task, targetConnections.get(i), produceEnd, consumerEnd, msg, RequestContextHolder.getRequestAttributes());
                executor.execute(writer);
            }
            // 6、关闭连接
            long start = System.currentTimeMillis();
            long now = System.currentTimeMillis();
            // 循环观测子线程是否存在异常，及时感知，否则发生异常时感知不及时，会等到全部结束或一小时才结束任务
            while ((now - start) / 1000 < CommonConstant.CONSUMER_TIME_OUT) {
                // 消费结束或存在异常，退出，否则继续观测
                if (consumerEnd.await(CommonConstant.LOOP_OBSERVATION_TIME_OUT, TimeUnit.SECONDS) || msg.isMarked()) {
                    break;
                }
                now = System.currentTimeMillis();
            }
            if (!msg.isMarked()) {
                log.info("DataTransferHandler end success");
                // 正常关闭线程池
                executor.shutdown();
            } else {
                log.warn("DataTransferHandler not end in {} seconds, shutdown now", CommonConstant.CONSUMER_TIME_OUT);
                // 关闭当前正在执行的子线程，只会interrupt正在执行的线程，需要在子线程判断终止标志
                executor.shutdownNow();
                throw new RuntimeException(msg.getReference());
            }
        } catch (Exception e) {
            log.error("DataTransferHandler error", e);
            throw e;
        } finally {
            closeConnections(targetConnections);
        }
    }

    private void executePreSQl(List<Connection> targetConnections, DataTransferTask task) throws SQLException {
        // 全量场景需要清空目标表
        if (!task.getIncrease()) {
            log.info("DataTransferHandler execute pre truncate sql");

            // 重命名旧表
            String tableBackUp = task.getTargetTable() + DateUtil.generateCurrentDayId(DateUtil.MMDDHHMM);
            String renameTableSql = getRenameTableSql(task, tableBackUp);
            try (PreparedStatement stmt = targetConnections.get(0).prepareStatement(renameTableSql)) {
                stmt.executeUpdate();
            } catch (Exception e) {
                log.error("truncate table error:{}", renameTableSql, e);
                throw new RuntimeException("truncate table error:" + renameTableSql + e.getMessage());
            }

            // 根据目标表创建新表，
            String createNewTableSql = getCreateNewTableSql(task, tableBackUp);
            try (PreparedStatement stmt = targetConnections.get(0).prepareStatement(createNewTableSql)) {
                stmt.executeUpdate();
            } catch (Exception e) {
                log.error("truncate table error:{}", createNewTableSql, e);
                throw new RuntimeException("truncate table error:" + createNewTableSql + e.getMessage());
            }
        }
    }

    private static String getRenameTableSql(DataTransferTask task, String tableBackUp) {
        String renameTableSql;
        switch (task.getTargetSource().getDriver()) {
            // 默认是mysql语法
            case CommonConstant.MYSQL_DRIVER:
            default:
                renameTableSql = String.format(CommonConstant.RENAME_TABLE, task.getTargetTable(), tableBackUp);
        }
        return renameTableSql;
    }

    private static String getCreateNewTableSql(DataTransferTask task, String tableBackUp) {
        String renameTableSql;
        switch (task.getTargetSource().getDriver()) {
            // 默认是mysql语法
            case CommonConstant.MYSQL_DRIVER:
            default:
                renameTableSql = String.format(CommonConstant.CREATE_TABLE_LIKE, task.getTargetTable(), tableBackUp);
        }
        return renameTableSql;
    }

    private void closeConnections(List<Connection> targetConnections) {
        for (Connection connection : targetConnections) {
            try {
                connection.close();
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }

    private void limitParam(DataTransferTask task) {
        // 部分值做限制
        if (task.getFetchSize() > CommonConstant.MAX_FETCH_SIZE || task.getFetchSize() < CommonConstant.NUMBER_ZERO) {
            task.setFetchSize(CommonConstant.MAX_FETCH_SIZE);
        }

        if (task.getConcurrencySize() > CommonConstant.MAX_CONCURRENCY_SIZE || task.getConcurrencySize() < CommonConstant.NUMBER_ZERO) {
            task.setConcurrencySize(CommonConstant.MAX_CONCURRENCY_SIZE);
        }

        if (task.getBatchInsertSize() > CommonConstant.MAX_BATCH_SIZE || task.getBatchInsertSize() < CommonConstant.NUMBER_ZERO) {
            task.setBatchInsertSize(CommonConstant.MAX_BATCH_SIZE);
        }
    }

    private Connection connect(DataSourceDO dataSourceDO) {
        try {
            Class.forName(dataSourceDO.getDriver());
            DriverManager.setLoginTimeout(CommonConstant.TIMEOUT_SECONDS);
            return DriverManager.getConnection(dataSourceDO.getJdbcUrl(), dataSourceDO.getUsername(), dataSourceDO.getPassword());
        } catch (Exception e) {
            throw new RuntimeException("initialize datasource connection error");
        }
    }


}